import os
import weakref
import numpy as np
from implement.variables.parameter import Parameter
from utils.common import to_tuple


# 层
# Layer类是持有参数并使用这些参数进行变换的类
class Layer:
    def __init__(self):
        self._params = set()

    # 设置实例变量时调用的特殊方法
    def __setattr__(self, name, value):
        """
        设置对象的属性。如果属性值为Parameter或Layer类的实例，将其添加到_params集合中。

        Parameters:
            name (str): 属性名。
            value (object): 属性值。

        Returns:
            None
        """
        if isinstance(value, (Parameter, Layer)):
            self._params.add(name)
        super().__setattr__(name, value)

    def __call__(self, *inputs):
        """
        通过调用forward方法计算函数的输出。同时，使用弱引用持有输入变量和输出变量。

        Parameters:
            *inputs: 输入变量，可以是多个变量。

        Returns:
            outputs: 函数计算得到的输出变量或变量组成的元组。
        """
        outputs = to_tuple(self.forward(*inputs))

        # 通过弱引用持有输入变量和输出变量
        self.inputs = [weakref.ref(x) for x in inputs]
        self.outputs = [weakref.ref(y) for y in to_tuple(outputs)]
        return outputs if len(outputs) > 1 else outputs[0]

    def forward(self, inputs):
        raise NotImplementedError

    def params(self):
        """
        返回模型的所有参数。

        Yields:
            Variable: 模型的参数变量。
        """

        for name in self._params:
            # yield使用方法和return相同，但yield是暂停处理并返回值，因此再次使用yield会恢复处理
            # 以上面的代码为例，每次调用params方法时暂停的处理都会重新运行。
            # 组合使用yield和for语句，即可按顺序取出参数。
            obj = self.__dict__[name]

            # 从Layer取出参数
            if isinstance(obj, Layer):
                # yield的函数叫作生成器，可以通过yield from来使用一个生成器创建另一个新的生成器
                yield from obj.params()
            else:
                # 如果是普通变量，则直接返回
                yield obj

    # 初始化导数
    def cleargrads(self):
        for param in self.params():
            param.cleargrad()

    def to_cpu(self):
        for param in self.params():
            param.to_cpu()

    def to_gpu(self):
        for param in self.params():
            param.to_gpu()

    def _flatten_params(self, params_dict=None, parent_key=""):
        """
        将模型的参数按照层级结构扁平化，并存储在params_dict中。

        Parameters:
            params_dict (dict, optional): 用于存储扁平化参数的字典，默认为一个空字典。
            parent_key (str, optional): 父级键名，用于构建参数在字典中的键名。

        Returns:
            None
        """
        for name in self._params:
            obj = self.__dict__[name]
            key = parent_key + '/' + name if parent_key else name

            if isinstance(obj, Layer):
                obj._flatten_params(params_dict, key)
            else:
                params_dict[key] = obj

    def save_weights(self, path: str) -> None:
        """
        将模型的权重保存到指定路径。

        Parameters:
            path (str): 保存权重的文件路径。

        Returns:
            None
        """
        self.to_cpu()

        params_dict = {}
        self._flatten_params(params_dict)
        array_dict = {key: param.data for key, param in params_dict.items()
                      if param is not None}
        try:
            np.savez_compressed(path, **array_dict)
        except (Exception, KeyboardInterrupt) as e:
            if os.path.exists(path):
                os.remove(path)
            raise

    def load_weights(self, path: str) -> None:
        """
        从指定路径加载模型的权重。

        Parameters:
            path (str): 权重文件的路径。

        Returns:
            None
        """
        try:
            # 尝试加载权重文件
            npz = np.load(path)
        except (FileNotFoundError, IOError):
            # 如果文件不存在或无法读取，抛出异常
            raise FileNotFoundError(f"Could not load weights from '{path}'. File not found or could not be read.")

        params_dict = {}
        self._flatten_params(params_dict)
        for key, param in params_dict.items():
            param.data = npz[key]